import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader


def get_train_loader(mean, std, batch_size=16, num_workers=2, shuffle=True):
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])

    cifar100_training = torchvision.datasets.CIFAR100(root='./data', train=True, download=True,
                                                      transform=transform_train)
    cifar100_training_loader = DataLoader(
        cifar100_training, shuffle=shuffle, num_workers=num_workers, batch_size=batch_size)

    return cifar100_training_loader


def get_test_loader(mean, std, batch_size=16, num_workers=2, shuffle=True):
    transform_test = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])
    cifar100_test = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)
    cifar100_test_loader = DataLoader(
        cifar100_test, shuffle=shuffle, num_workers=num_workers, batch_size=batch_size)

    return cifar100_test_loader

if __name__ == '__main__':
    mean = [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]
    std = [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]

    def denormalize(imgs, mean, std):
        if not torch.is_tensor(imgs):
            raise TypeError("Input should be a torch tensor.")
        for t, m, s in zip(imgs, mean, std):
            t.mul_(s).add_(m)
        return imgs

    import matplotlib.pyplot as plt
    from torchvision.utils import make_grid

    def show_imgs(images, labels):
        import matplotlib
        matplotlib.use('TkAgg')
        images = denormalize(images, mean, std)
        img_grid = make_grid(images, nrow=4, padding=10, normalize=True)
        plt.imshow(img_grid.permute(1, 2, 0))
        plt.title(f"Labels: {labels}")
        plt.show()

    test_loader = get_test_loader(mean, std, batch_size=16, num_workers=2, shuffle=False)

    imgs, labels = next(iter(test_loader))
    show_imgs(imgs, labels)
    print(labels)
